from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.exceptions import InvalidToken, AuthenticationFailed
from rest_framework_simplejwt.tokens import AccessToken
from django.conf import settings
from .models import UserDevice

class CustomJWTAuthentication(JWTAuthentication):
    def get_validated_token(self, raw_token):
        """验证 token 并检查设备状态"""
        # 首先验证 token 本身（包括过期检查）
        validated_token = super().get_validated_token(raw_token)
        
        try:
            # 获取设备信息
            device_id = validated_token.get('device_id')
            user_id = validated_token.get('user_id')
            token_version = validated_token.get('token_version')
            
            if device_id and user_id:
                # 检查设备状态和 token 版本
                device = UserDevice.objects.filter(
                    device_id=device_id,
                    user_id=user_id
                ).first()
                
                if device:
                    if not device.is_active:
                        raise InvalidToken({
                            'detail': 'Device has been logged out',
                            'code': 'token_not_valid'
                        })
                    
                    if token_version != device.token_version:
                        raise InvalidToken({
                            'detail': 'Token version mismatch',
                            'code': 'token_not_valid'
                        })
                else:
                    raise InvalidToken({
                        'detail': 'Device not found',
                        'code': 'token_not_valid'
                    })
            
            return validated_token
            
        except Exception as e:
            raise InvalidToken({
                'detail': str(e),
                'code': 'token_not_valid'
            })
